package org.gjw.config;

import cn.hutool.core.util.StrUtil;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import lombok.NoArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.gjw.bean.DataSourceInfo;
import org.gjw.mapper.DataSourceInfoMapper;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.jdbc.DataSourceBuilder;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import javax.sql.DataSource;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author guojunwang
 * @date 2022-05-10 21:42
 */
@Component
@NoArgsConstructor
@Slf4j
public class DataSourceHolder {

    @Autowired
    private AbstractRoutingDataSource routingDataSource;

    @Resource
    private DataSourceInfoMapper dataSourceInfoMapper;

    /**
     * 存放所有数据源,其实 key = dataSourceKey , value = DataSource
     */
    protected final ConcurrentHashMap<Object, Object> dataSourceMap = new ConcurrentHashMap<>(64);

    /**
     * 用于初始化数据源时加锁
     */
    private final ConcurrentHashMap<String, Object> dataSourceInitLock = new ConcurrentHashMap<>(64);

    /**
     * 保存当前的 数据源key
     */
    private static ThreadLocal<String> currDataSourceKey = new ThreadLocal<>();

    /**
     * 是否使用默认的数据源
     */
    private static ThreadLocal<Boolean> useDefaultDs = new ThreadLocal<>();

    public static void setDataSourceKey(String dataSourceKey){
        log.info("数据源切换:{} -> {}", currDataSourceKey.get(), dataSourceKey  );
        currDataSourceKey.set(dataSourceKey);
        unUseDefaultDataSource();
    }

    public static void useDefaultDataSource(){
        useDefaultDs.set(true);
    }
    public static void unUseDefaultDataSource(){
        useDefaultDs.remove();
    }

    /**
     * 获取当前的数据源 key
     * @return 当前 dataSourcekey
     */
    public String getDataSourceKey(){

        //获取指定的数据源 KEY
        final String currentDsKey = currDataSourceKey.get();

        //如果指定使用默认数据源 或者 未指定数据源，返回null
        if( Objects.equals(useDefaultDs.get(),true) || StrUtil.isBlank( currentDsKey ) ){
            log.info("采用默认数据源");
            return null;
        }

        //如果指定的数据源已经被初始化
        if( dataSourceMap.containsKey(currentDsKey) ){
            log.info("采用动态数据源,dataSourceKey:{}",currentDsKey);
            return currentDsKey;
        }

        //如果数据源未被初始化，尝试初始化数据源
        return initDataSource( currentDsKey );
    }


    /**
     * 初始化数据源
     * @return 返回数据源key
     */
    private String initDataSource( final String dataSourceKey ) {

        //如果不存在，设置一个对象，主要用于加锁
        Object lockObj = dataSourceInitLock.computeIfAbsent(dataSourceKey, key -> new Object());


        synchronized ( lockObj ){
            log.info("开始初始化数据源,dataSourceKey:{}",dataSourceKey);

            if( dataSourceMap.containsKey(dataSourceKey) ){
                log.info("数据源已被其他线程初始化,dataSourceKey:{}",dataSourceKey);
                return dataSourceKey;
            }

            try {
                //使用默认的数据源
                useDefaultDataSource();

                //根据dataSourceKey,查询数据库配置信息
                LambdaQueryWrapper<DataSourceInfo> wrapper = Wrappers.lambdaQuery();
                wrapper.eq(DataSourceInfo::getDataSourceKey,dataSourceKey);
                wrapper.last( "limit 1");

                DataSourceInfo dataSourceInfo = dataSourceInfoMapper.selectOne(wrapper);

                if( dataSourceInfo == null ){
                    //这里要抛异常，否则会使用默认的数据源
                    log.error("加载数据源失败，数据源配置信息不存在，dataSourceKey:{}",dataSourceKey);
                    return null;
                }

                Class dataSourceType = Class.forName( dataSourceInfo.getDataSourceType() );

                //创建数据源
                DataSource dataSource = DataSourceBuilder.create()
                    .driverClassName(dataSourceInfo.getDriverClassName())
                    .username(dataSourceInfo.getUsername())
                    .password(dataSourceInfo.getPassword())
                    .url(dataSourceInfo.getUrl())
                    .type(dataSourceType)
                    .build();

                //放入数据源map
                dataSourceMap.put(dataSourceKey,dataSource);
                //数据源改变后需要刷新
                refreshDataSource();

                //创建成功后返回
                return dataSourceKey;
            }catch( ClassNotFoundException ignored){
            } finally {
                unUseDefaultDataSource();
            }

        }
        return null;
    }

    /**
     * 删除数据源
     * @param dataSourceKey 数据源Key
     */
    public void deleteDataSource( String dataSourceKey ){
        dataSourceMap.remove(dataSourceKey);
        refreshDataSource();
    }


    /**
     * 刷新数据源
     */
    public void refreshDataSource(){
        log.info("刷新动态数据源");
        routingDataSource.afterPropertiesSet();
    }

}
